from flask import Blueprint, request, jsonify, send_file, Response
import google.generativeai as genai
import os
import json
from backend.app.utils.excel import generate_character_excel

bp = Blueprint('chat', __name__)

# 存储API密钥的全局变量
GEMINI_API_KEY = None
current_model = None

# 确保config目录存在
config_dir = os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(__file__))), 'config')
if not os.path.exists(config_dir):
    os.makedirs(config_dir)

HISTORY_FILE = os.path.join(config_dir, 'chat_history.json')
CONFIG_FILE = os.path.join(config_dir, 'config.json')


def load_chat_history():
    try:
        if os.path.exists(HISTORY_FILE):
            with open(HISTORY_FILE, 'r', encoding='utf-8') as f:
                return json.load(f)
    except Exception as e:
        print(f"Error loading chat history: {e}")
    return []


def save_chat_history(history):
    try:
        # 确保历史记录是一个列表
        if not isinstance(history, list):
            history = []

        # 直接保存新的历史记录，不进行合并
        with open(HISTORY_FILE, 'w', encoding='utf-8') as f:
            json.dump(history, f, ensure_ascii=False, indent=2)
    except Exception as e:
        print(f"Error saving chat history: {e}")
        raise e


def load_config():
    try:
        if os.path.exists(CONFIG_FILE):
            with open(CONFIG_FILE, 'r', encoding='utf-8') as f:
                config = json.load(f)
                if not isinstance(config, dict):
                    print("配置文件格式错误：不是有效的JSON对象")
                    return {'api_key': '', 'current_model': ''}
                config.setdefault('api_key', '')
                config.setdefault('current_model', '')
                return config
        else:
            print(f"配置文件不存在：{CONFIG_FILE}")
    except Exception as e:
        print(f"配置加载错误: {e}")
    return {'api_key': '', 'current_model': ''}


def save_config(config):
    try:
        with open(CONFIG_FILE, 'w', encoding='utf-8') as f:
            json.dump(config, f, ensure_ascii=False, indent=2)
    except Exception as e:
        print(f"配置保存错误: {e}")


def init_api_key():
    config = load_config()
    api_key = config.get('api_key', '').strip()
    current_model_name = config.get('current_model', '').strip()

    if not api_key:
        print("未找到有效的API Key配置")
        return False

    try:
        genai.configure(api_key=api_key, transport='rest')
        models = genai.list_models()
        if not any('generateContent' in model.supported_generation_methods for model in models):
            print("API Key验证失败：未找到可用的生成模型")
            return False

        global GEMINI_API_KEY, current_model
        GEMINI_API_KEY = api_key
        current_model = current_model_name if current_model_name else ''

        print("API Key初始化成功")
        return True
    except Exception as e:
        print(f"API Key初始化失败: {str(e)}")
        return False


# 在应用启动时初始化API Key
init_api_key()


@bp.route('/validate_key', methods=['POST'])
def validate_key():
    try:
        data = request.json
        api_key = data.get('apiKey', '').strip()

        if not api_key:
            return jsonify({'data': {'success': False, 'error': 'API Key不能为空'}}), 400

        genai.configure(api_key=api_key, transport='rest')
        global GEMINI_API_KEY
        GEMINI_API_KEY = api_key

        models = genai.list_models()
        available_models = [
            {
                'name': model.name,
                'displayName': model.name.replace('models/', ''),
                'supported_methods': model.supported_generation_methods
            }
            for model in models
            if 'generateContent' in model.supported_generation_methods
        ]

        if not available_models:
            return jsonify({'data': {'success': False, 'error': '没有找到可用的模型'}}), 400

        config = load_config()
        config['api_key'] = api_key
        save_config(config)

        current_model_name = config.get('current_model', '')
        return jsonify({
            'data': {
                'success': True,
                'models': available_models,
                'current_model': current_model_name
            }
        })

    except Exception as e:
        return jsonify({'data': {'success': False, 'error': f'验证过程出错：{str(e)}'}}), 500


@bp.route('/set_model', methods=['POST'])
def set_model():
    try:
        data = request.json
        model_name = data.get('model', '').strip()
        if not model_name:
            return jsonify({'data': {'error': '模型名称不能为空'}}), 400

        config = load_config()
        config['current_model'] = model_name
        save_config(config)

        global current_model
        current_model = model_name
        return jsonify({'data': {'message': '模型设置成功'}})
    except Exception as e:
        return jsonify({'data': {'error': f'设置模型失败：{str(e)}'}}), 500


@bp.route('/config', methods=['GET'])
def get_config():
    try:
        config = load_config()
        return jsonify({
            'data': {
                'success': True,
                'config': config
            }
        })
    except Exception as e:
        return jsonify({'data': {'success': False, 'error': str(e)}}), 400


@bp.route('/history', methods=['GET'])
def get_history():
    try:
        history = load_chat_history()
        return jsonify({
            'data': {
                'success': True,
                'history': history
            }
        })
    except Exception as e:
        return jsonify({'data': {'success': False, 'error': str(e)}}), 400


@bp.route('/clear_history', methods=['POST'])
def clear_history():
    try:
        # 清空历史记录文件
        save_chat_history([])
        return jsonify({
            'data': {
                'success': True,
                'message': '历史记录已清空'
            }
        })
    except Exception as e:
        return jsonify({
            'data': {
                'success': False,
                'error': f'清空历史记录失败：{str(e)}'
            }
        }), 500


@bp.route('/revoke_message', methods=['POST'])
def revoke_message():
    try:
        data = request.json
        index = data.get('index')
        history = load_chat_history()

        if not history or index < 0 or index >= len(history):
            return jsonify({
                'data': {
                    'success': False,
                    'error': '无效的消息索引'
                }
            }), 400

        # 如果要撤回的是用户消息，同时也要删除对应的AI回复
        if index < len(history) - 1 and history[index]['role'] == 'user':
            history.pop(index + 1)  # 删除AI回复
        history.pop(index)  # 删除当前消息

        save_chat_history(history)
        return jsonify({
            'data': {
                'success': True,
                'message': '消息已撤回',
                'history': history
            }
        })
    except Exception as e:
        return jsonify({
            'data': {
                'success': False,
                'error': f'撤回消息失败：{str(e)}'
            }
        }), 500


@bp.route('/chat', methods=['GET', 'POST'])
def chat():
    try:
        data = request.json
        message = data.get('message', '').strip()
        model_name = data.get('model', '').strip()
        stream_mode = data.get('stream', True)  # 默认使用流式响应
        history = load_chat_history()

        if not message:
            return jsonify({'data': {'success': False, 'error': '消息不能为空'}}), 400

        if not model_name:
            return jsonify({'data': {'success': False, 'error': '请选择模型'}}), 400

        config = load_config()
        api_key = config.get('api_key', '')

        if not api_key:
            return jsonify({'data': {'success': False, 'error': '请先设置API Key'}}), 400

        genai.configure(api_key=api_key, transport='rest')

        try:
            model = genai.GenerativeModel(model_name)
            chat = model.start_chat()

            # 非流式响应模式
            if not stream_mode:
                try:
                    response = chat.send_message(message, stream=False)
                    response_text = response.text

                    # 保存完整的对话历史
                    history.append({'role': 'user', 'content': message})
                    history.append({'role': 'assistant', 'content': response_text})
                    save_chat_history(history)

                    return jsonify({
                        'content': response_text,
                        'done': True
                    })
                except Exception as e:
                    error_msg = f"生成回复失败：{str(e)}"
                    return jsonify({'error': error_msg}), 400

            # 流式响应模式
            def generate():
                try:
                    response = chat.send_message(message, stream=True)
                    accumulated_text = ""

                    for chunk in response:
                        if chunk.text:
                            accumulated_text += chunk.text
                            yield f"data: {json.dumps({'content': chunk.text, 'done': False})}\n\n"

                    # 保存完整的对话历史
                    history.append({'role': 'user', 'content': message})
                    history.append({'role': 'assistant', 'content': accumulated_text})
                    save_chat_history(history)

                    # 发送完成标记
                    yield f"data: {json.dumps({'content': '', 'done': True})}\n\n"

                except Exception as e:
                    error_msg = f"生成回复失败：{str(e)}"
                    if "Unknown error trying to retrieve streaming response" in str(e):
                        error_msg = "流式响应失败，请尝试使用非流式模式 (stream=False)"
                    yield f"data: {json.dumps({'error': error_msg})}\n\n"

            response = Response(generate(), mimetype='text/event-stream')
            response.headers['Cache-Control'] = 'no-cache'
            response.headers['X-Accel-Buffering'] = 'no'
            response.headers['Access-Control-Allow-Origin'] = '*'
            return response

        except Exception as e:
            return jsonify({'data': {'success': False, 'error': f'生成回复失败：{str(e)}'}}), 400

    except Exception as e:
        return jsonify({'data': {'success': False, 'error': f'请求处理失败：{str(e)}'}}), 400


@bp.route('/save_history', methods=['POST'])
def save_history_endpoint():
    try:
        data = request.json
        history = data.get('history', [])

        if not isinstance(history, list):
            return jsonify({'data': {'error': '历史记录格式错误'}}), 400

        save_chat_history(history)
        return jsonify({'data': {'message': '保存成功'}})
    except Exception as e:
        return jsonify({'data': {'error': f'保存失败：{str(e)}'}}), 400


@bp.route('/generate_excel', methods=['POST'])
def generate_excel_endpoint():
    try:
        data = request.json
        text = data.get('content', '')
        headers = data.get('headers', ['镜头编号', '镜头描述', '风格', '角色', '场景', '中文描述', 'AI 分镜词'])

        if not text:
            return jsonify({'data': {'error': '内容不能为空'}}), 400

        result = generate_character_excel(text, headers)

        if result['success']:
            return jsonify({
                'data': {
                    'file_name': result['file_name'],
                    'message': result['message'],
                    'download_url': f'/download/{result["file_name"]}'
                }
            })
        else:
            return jsonify({'data': result}), 400

    except Exception as e:
        return jsonify({'data': {'error': f'请求处理失败：{str(e)}'}}), 400


@bp.route('/download/<filename>')
def download_file(filename):
    try:
        return send_file(
            os.path.join(os.getcwd(), filename),
            as_attachment=True,
            download_name=filename
        )
    except Exception as e:
        return jsonify({'data': {'error': f'文件下载失败：{str(e)}'}}), 400


@bp.route('/delete_message', methods=['POST'])
def delete_message():
    try:
        data = request.json
        index = data.get('index')

        if os.path.exists(HISTORY_FILE):
            with open(HISTORY_FILE, 'r', encoding='utf-8') as f:
                history = json.load(f)
        else:
            history = []

        if 0 <= index < len(history):
            history.pop(index)

            with open(HISTORY_FILE, 'w', encoding='utf-8') as f:
                json.dump(history, f, ensure_ascii=False, indent=2)

            return jsonify({
                'data': {
                    'message': '消息已删除'
                }
            })
        else:
            return jsonify({
                'data': {
                    'error': '无效的消息索引'
                }
            })

    except Exception as e:
        return jsonify({
            'data': {
                'error': f'删除失败：{str(e)}'
            }
        }), 500


@bp.route('/chat_storyboard', methods=['POST'])
def chat_storyboard():
    try:
        config = load_config()
        data = request.json
        message = data.get('message', '').strip()
        model_name = data.get('model', data.get('model_name', '') if data else '').strip()
        if not message:
            return jsonify({'data': {'success': False, 'error': '消息不能为空'}}), 400
        if not model_name:
            return jsonify({'data': {'success': False, 'error': '请选择模型'}}), 400
        api_key = config.get('api_key', '')
        if not api_key:
            return jsonify({'data': {'success': False, 'error': '请先设置API Key'}}), 400
        genai.configure(api_key=api_key, transport='rest')
        try:
            model = genai.GenerativeModel(model_name)
            chat = model.start_chat()

            def generate():
                try:
                    response = chat.send_message(message, stream=True)
                    accumulated_text = ""

                    for chunk in response:
                        print(chunk)
                        if chunk.text:
                            accumulated_text += chunk.text
                            yield f"data: {json.dumps({'content': chunk.text, 'done': False})}\n\n"
                    # 发送完成标记
                    yield f"data: {json.dumps({'content': '', 'done': True})}\n\n"
                except Exception as e:
                    error_msg = f"生成回复失败：{str(e)}"
                    yield f"data: {json.dumps({'error': error_msg})}\n\n"

            response = Response(generate(), mimetype='text/event-stream')
            response.headers['Cache-Control'] = 'no-cache'
            response.headers['X-Accel-Buffering'] = 'no'
            response.headers['Access-Control-Allow-Origin'] = '*'
            return response
        except Exception as e:
            return jsonify({'data': {'success': False, 'error': f'生成回复失败：{str(e)}'}}), 400
    except Exception as e:
        return jsonify({'data': {'success': False, 'error': f'请求处理失败：{str(e)}'}}), 400
